/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "kernel_builder.h"

#include "engine/base_engine.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/utils/tensor_utils.h"
#include "graph/utils/type_utils.h"
#include "securec.h"
#include "util/constant.h"
#include "util/log.h"
#include "util/util.h"

using namespace std;
using namespace ge;
using namespace ::aicpu::FWKAdapter;

namespace {
const string kOutputMemsizeVector = "output_memsize_vector";
const string kOutputAlignmentVector = "output_alignment_vector";
const string kAicpuEngine = "DNN_VM_AICPU_ASCEND";
const string kAllShape = "_AllShape";
const uint64_t kAlignmentValue = 64;
}  // namespace

namespace aicpu {
// calculate and set outputs size
Status KernelBuilder::SetOutPutsSize(shared_ptr<OpDesc> &op_desc_ptr) const {
  AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, INPUT_PARAM_NULL)
  int32_t shape_type = 0;
  if (AttrUtils::HasAttr(op_desc_ptr, kAttrNameUnknownShape)) {
    // unknow shape
    CHECK_RES_BOOL(
        AttrUtils::GetInt(op_desc_ptr, ATTR_NAME_UNKNOWN_SHAPE_TYPE, shape_type),
        INVOKE_GRAPH_ITF_FAILED,
        AICPU_REPORT_CALL_ERROR(
            "Call ge::AttrUtils::GetInt failed to get attr[%s], op[%s].",
            ATTR_NAME_UNKNOWN_SHAPE_TYPE.c_str(), op_desc_ptr->GetName().c_str()))
    if (shape_type == DEPEND_COMPUTE) {
      // unknow type 4, output is ResultSummary
      AICPU_CHECK_RES_WITH_LOG(SetOutSizeForSummary(op_desc_ptr),
          "Call SetOutSizeForSummary function failed, op[%s].",
          op_desc_ptr->GetName().c_str())
      return SUCCESS;
    }
  }

  vector<int64_t> output_memsize_vec;
  vector<int64_t> output_alignment_vec;
  size_t output_size = op_desc_ptr->GetOutputsSize();
  for (size_t i = 0; i < output_size; i++) {
    int64_t output_mem_size = 0;
    GeTensorDesc output_tensor = op_desc_ptr->GetOutputDesc(i);
    DataType output_data_type = output_tensor.GetDataType();
    if (shape_type == DEPEND_SHAPE_RANGE) {
      vector<pair<int64_t, int64_t>> shape_range;
      // try get shape range from tensor desc
      AICPU_CHECK_RES_WITH_LOG(output_tensor.GetShapeRange(shape_range),
          "Call GetShapeRange function failed for %zuth output, op[%s].",
          i, op_desc_ptr->GetName().c_str())
      if (!shape_range.empty()) {
        // unknow type3, calc outputs size by ShapeRange
        AICPU_CHECK_RES_WITH_LOG(
            GetOutSizeByShapeRange(output_data_type, shape_range, output_mem_size),
            "Call GetOutSizeByShapeRange function failed for %zuth output, op[%s].",
            i, op_desc_ptr->GetName().c_str())
      } else {
        const GeShape &output_shape = output_tensor.GetShape();
        aicpu::State state = GetTotalSizeByShapeAndType(
            output_data_type, output_shape, output_mem_size);
        AICPU_IF_BOOL_EXEC(state.state != ge::SUCCESS,
            AICPU_REPORT_INNER_ERROR("Overflow occurred when calculate total "
                "bytes of output[%zu], %s, op[%s]", i, state.msg.c_str(),
                op_desc_ptr->GetName().c_str());
            return state.state)
      }
    } else {
      const GeShape &output_shape = output_tensor.GetShape();
      aicpu::State state = GetTotalSizeByShapeAndType(
          output_data_type, output_shape, output_mem_size);
      AICPU_IF_BOOL_EXEC(state.state != ge::SUCCESS,
          AICPU_REPORT_INNER_ERROR("Overflow occurred when calculate total "
              "bytes of output[%zu], %s, op[%s]", i, state.msg.c_str(),
              op_desc_ptr->GetName().c_str());
          return state.state)
    }

    AICPUE_LOGI("After the output [%lu]'s calc, the total size is [%lld]", i,
                output_mem_size);
    if (output_mem_size > LONG_MAX) {
      AICPU_REPORT_INNER_ERROR("Overflow occured when calculate %zuth output "
          "total bytes, data type[%s], shape[%s], op[%s].", i,
          ge::TypeUtils::DataTypeToSerialString(output_data_type).c_str(),
          DebugString(output_tensor.GetShape().GetDims()).c_str(),
          op_desc_ptr->GetName().c_str());
      return DATA_OVERFLOW;
    }
    TensorUtils::SetSize(output_tensor, output_mem_size);
    (void)output_memsize_vec.emplace_back(output_mem_size);
    (void)output_alignment_vec.emplace_back(kAlignmentValue);
    AICPU_CHECK_RES(op_desc_ptr->UpdateOutputDesc(i, output_tensor));
  }

  CHECK_RES_BOOL(
      AttrUtils::SetListInt(op_desc_ptr, kOutputMemsizeVector, output_memsize_vec),
      INVOKE_GRAPH_ITF_FAILED,
      AICPU_REPORT_CALL_ERROR(
          "Call ge::AttrUtils::SetListInt failed to set attr[%s], op[%s]",
          kOutputMemsizeVector.c_str(), op_desc_ptr->GetName().c_str()))
  CHECK_RES_BOOL(AttrUtils::SetListInt(op_desc_ptr, kOutputAlignmentVector,
                    output_alignment_vec),
      INVOKE_GRAPH_ITF_FAILED,
      AICPU_REPORT_CALL_ERROR(
          "Call ge::AttrUtils::SetListInt failed to set attr[%s], op[%s]",
          kOutputAlignmentVector.c_str(), op_desc_ptr->GetName().c_str()))
  return SUCCESS;
}

// Set outputs size for unknow type 4 which output is result summary
Status KernelBuilder::SetOutSizeForSummary(
    shared_ptr<OpDesc> &op_desc_ptr) const {
  vector<int64_t> output_memsize_vec;
  vector<int64_t> output_alignment_vec;

  size_t output_size = op_desc_ptr->GetOutputsSize();
  for (size_t i = 0; i < output_size; i++) {
    // output is ResultSummary, 4 means number of fields for ResultSummary
    int64_t output_mem_size = 4 * sizeof(uint64_t);
    GeTensorDesc output_tensor = op_desc_ptr->GetOutputDesc(i);
    TensorUtils::SetSize(output_tensor, static_cast<uint32_t>(output_mem_size));
    (void)output_memsize_vec.emplace_back(output_mem_size);
    (void)output_alignment_vec.emplace_back(kAlignmentValue);
    AICPU_CHECK_RES(op_desc_ptr->UpdateOutputDesc(i, output_tensor));
  }
  CHECK_RES_BOOL(
      AttrUtils::SetListInt(op_desc_ptr, kOutputMemsizeVector, output_memsize_vec),
      INVOKE_GRAPH_ITF_FAILED,
      AICPU_REPORT_CALL_ERROR(
          "Call ge::AttrUtils::SetListInt failed to set attr[%s], op[%s]",
          kOutputMemsizeVector.c_str(), op_desc_ptr->GetName().c_str()))
  CHECK_RES_BOOL(AttrUtils::SetListInt(op_desc_ptr, kOutputAlignmentVector,
                     output_alignment_vec),
    INVOKE_GRAPH_ITF_FAILED,
    AICPU_REPORT_CALL_ERROR(
        "Call ge::AttrUtils::SetListInt failed to set attr[%s], op[%s]",
        kOutputAlignmentVector.c_str(), op_desc_ptr->GetName().c_str()))

  AICPUE_LOGI(
      "Set output size with ResultSummary for unknow type 4 ok, op: %s.",
      op_desc_ptr->GetName().c_str());
  return SUCCESS;
}

State KernelBuilder::GetTotalSizeByShapeAndType(const DataType &data_type,
                                                const GeShape &ge_shape,
                                                int64_t &total_size) const {
  int32_t data_size = GetDataTypeSize(data_type);
  AICPU_CHECK_GREATER_THAN_ZERO(data_size, DATA_TYPE_UNDEFILED,
      "Invalid data type[%s].",
      ge::TypeUtils::DataTypeToSerialString(data_type).c_str())
  vector<int64_t> dims = ge_shape.GetDims();
  int64_t total_num = 1;
  size_t dim_size = dims.size();
  for (size_t i = 0; i < dim_size; i++) {
    if (dims[i] == UNKNOWN_DIM || dims[i] == UNKNOWN_DIM_NUM) {
      AICPUE_LOGW("Unknow output shape [%d], skip it.", dims[i]);
      break;
    }
    if (dims[i] < 0) {
      string err_msg =
          Stringcat("dim[", i, "] is invalid, shape is ", DebugString(dims));
      aicpu::State state(GE_SHAPE_SIZE_INVAILD, err_msg);
      return state;
    }
    CHECK_INT64_MUL_OVERFLOW(total_num, dims[i], aicpu::State(DATA_OVERFLOW),
        "Overflow when calculate tensor total bytes, shape[%s], data type[%s],",
        DebugString(dims).c_str(),
        ge::TypeUtils::DataTypeToSerialString(data_type).c_str())
    total_num *= dims[i];
  }

  CHECK_INT64_MUL_OVERFLOW(total_num, data_size, aicpu::State(DATA_OVERFLOW),
      "Overflow when calculate tensor total bytes, shape[%s], data type[%s],",
      DebugString(dims).c_str(),
      ge::TypeUtils::DataTypeToSerialString(data_type).c_str())
  total_size = total_num * data_size;
  return aicpu::State(SUCCESS);
}

Status KernelBuilder::GetWorkspaceInfo(const OpDescPtr &op_desc_ptr,
                                       uint8_t *data_mem_base,
                                       uint64_t &workspace_bytes_size) const {
  // Get the workspace's mem base and mem size
  AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, INPUT_PARAM_NULL)

  uintptr_t data_mem_base_addr = reinterpret_cast<uintptr_t>(data_mem_base);
  int64_t workspace_mem_base = static_cast<int64_t>(data_mem_base_addr);

  // Get workspace total size, the data source is from calcOpRunningParam
  vector<int64_t> workspace_bytes = op_desc_ptr->GetWorkspaceBytes();
  AICPUE_LOGI("The workspace_bytes vector's size is [%lu], op[%s].",
              workspace_bytes.size(), op_desc_ptr->GetName().c_str());
  CHECK_EQUAL(workspace_bytes.size(), 1, INPUT_PARAM_VALID,
      "workspace bytes vector's size[%zu] should be 1, op[%s].",
      workspace_bytes.size(), op_desc_ptr->GetName().c_str())
  CHECK_RES_BOOL((workspace_bytes[0] >= 0), INPUT_PARAM_VALID,
      AICPU_REPORT_INNER_ERROR("Invalied workspace bytes[%ld], op[%s]",
          workspace_bytes[0], op_desc_ptr->GetName().c_str()))
  workspace_bytes_size = static_cast<uint64_t>(workspace_bytes[0]);
  AICPUE_LOGI("The workspace byte size is [%lld], op[%s]", workspace_bytes_size,
              op_desc_ptr->GetName().c_str());
  // Check workspace mem
  CHECK_INT64_ADD_OVERFLOW(workspace_mem_base, workspace_bytes_size, DATA_OVERFLOW,
      "Overflow when offset workspace bytes[%lu] base on workspace addr[%ld], op[%s]",
          workspace_bytes_size, workspace_mem_base, op_desc_ptr->GetName().c_str())

  return SUCCESS;
}

// calculate output size for unknow type 3 which output shape is range
Status KernelBuilder::GetOutSizeByShapeRange(
    const DataType &data_type, const vector<pair<int64_t, int64_t>> &shape_range,
    int64_t &total_size) const {
  int32_t data_size = GetDataTypeSize(data_type);
  AICPU_CHECK_GREATER_THAN_ZERO(data_size, DATA_TYPE_UNDEFILED,
      "Invalid data type[%s].",
      ge::TypeUtils::DataTypeToSerialString(data_type).c_str())

  int64_t total_num = 1;
  for (size_t index = 0; index < shape_range.size(); ++index) {
    const auto &dim_item = shape_range[index];
    // unknow shape
    if (dim_item.second == UNKNOWN_DIM || dim_item.second == UNKNOWN_DIM_NUM) {
      AICPUE_LOGW("Unknow output shape [%d], skip it.", dim_item.second);
      break;
    }
    AICPU_CHECK_GREATER_THAN_ZERO(dim_item.second, GE_SHAPE_SIZE_INVAILD,
        "Invalid value[%ld] of dim[%zu].", dim_item.second, index)
    CHECK_INT64_MUL_OVERFLOW(total_num, dim_item.second, DATA_OVERFLOW,
        "Overflow when calculate tensor shape range total bytes")
    total_num *= dim_item.second;
  }
  CHECK_INT64_MUL_OVERFLOW(total_num, data_size, DATA_OVERFLOW,
      "Overflow when calculate tensor shape range total bytes")
  total_size = total_num * data_size;
  AICPUE_LOGI("Calc output size by shape range ok, total_size[%ld].", total_size);
  return SUCCESS;
}

ge::Status KernelBuilder::MakeExtInfoForOpName(
    const ge::OpDescPtr &op_desc_ptr, std::vector<char> &task_ext_info) const {
  AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, INPUT_PARAM_NULL)
  // WARNING: OP NAME MUST BE THE FIRST EXTEND INFO FOR RUNTIME!!!
  // make extend info for op name
  string op_name = op_desc_ptr->GetName();
  // calc ExtInfo length
  uint64_t extend_info_len = task_ext_info.size();
  // ext info: op name, value: op_name.length()
  extend_info_len += (kExtInfoHeadSize + op_name.length());

  uint64_t ext_info_offset = task_ext_info.size();
  task_ext_info.resize(extend_info_len, 0);
  char *ext_info_buf = task_ext_info.data();
  AICPU_CHECK_NOTNULL_ERRCODE(ext_info_buf, INPUT_PARAM_NULL)
  // create ExtInfo struct
  AICPU_CHECK_NOTNULL(ext_info_buf)
  ExtInfo *ext_info = reinterpret_cast<ExtInfo *>(
      reinterpret_cast<uintptr_t>(ext_info_buf) + ext_info_offset);
  ext_info->infoType = FWK_ADPT_EXT_OP_NAME;
  ext_info->infoLen = op_name.length();
  ext_info_offset += kExtInfoHeadSize;
  char *op_name_buf = reinterpret_cast<char *>(
      reinterpret_cast<uintptr_t>(ext_info_buf) + ext_info_offset);
  errno_t cpy_ret =
      memcpy_s(op_name_buf, ext_info->infoLen, op_name.c_str(), op_name.length());
  if (cpy_ret != EOK) {
    AICPUE_LOGW(
        "Generate extend info for op[%s] failed, ext info len[%u], ret[%d].",
        op_name.c_str(), ext_info->infoLen, cpy_ret);
  }
  AICPUE_LOGI("Make ext_info for op[%s], ext info len[%lu].", op_name.c_str(),
              ext_info->infoLen);
  return SUCCESS;
}

// Make and init task extend info
Status KernelBuilder::MakeBaseExtInfo(const OpDescPtr &op_desc_ptr,
                                      vector<char> &task_ext_info) const {
  AICPU_CHECK_NOTNULL_ERRCODE(op_desc_ptr, INPUT_PARAM_NULL)
  // set attr _all_shape
  CHECK_RES_BOOL(
      AttrUtils::SetBool(op_desc_ptr, kAllShape, true), INVOKE_GRAPH_ITF_FAILED,
      AICPU_REPORT_CALL_ERROR(
          "Call ge::AttrUtils::SetBool failed to set attr[%s], op[%s].",
          kAllShape.c_str(), op_desc_ptr->GetName().c_str()))

  // calc ExtInfo length
  uint64_t extend_info_len = task_ext_info.size();
  // ext info 1: unknown shape_type, value: int32
  extend_info_len += (kExtInfoHeadSize + sizeof(int32_t));
  // ext info 2: bitmap, value: uint64
  extend_info_len += (kExtInfoHeadSize + sizeof(uint64_t));
  // ext info 3: topic type, value: enum
  extend_info_len += (kExtInfoHeadSize + sizeof(FWKExtTopicType));
  // ext info 4: input ShapeAndType, value: input_num * sizeof(ShapeAndType)
  size_t input_num = op_desc_ptr->GetInputsSize();
  size_t output_num = op_desc_ptr->GetOutputsSize();
  extend_info_len += kExtInfoHeadSize;
  extend_info_len += (input_num * sizeof(ShapeAndType));
  // ext info 5: output ShapeAndType, value: output_num * sizeof(ShapeAndType)
  extend_info_len += kExtInfoHeadSize;
  extend_info_len += (output_num * sizeof(ShapeAndType));

  uint64_t ext_info_offset = task_ext_info.size();
  task_ext_info.resize(extend_info_len, 0);
  char *ext_info_buf = task_ext_info.data();
  AICPU_CHECK_NOTNULL_ERRCODE(ext_info_buf, INPUT_PARAM_NULL)
  // initialize extend info
  // init ext info 1: unknown shape_type
  int32_t unknow_shape_type = 0;
  if (!AttrUtils::GetInt(op_desc_ptr, ATTR_NAME_UNKNOWN_SHAPE_TYPE,
                         unknow_shape_type)) {
    AICPU_REPORT_CALL_ERROR(
        "Call ge::AttrUtils::GetInt failed to get attr[%s], op[%s].",
        ATTR_NAME_UNKNOWN_SHAPE_TYPE.c_str(), op_desc_ptr->GetName().c_str());
    return GET_ATTR_FAILED;
  }
  ExtInfo *ext_info = reinterpret_cast<ExtInfo *>(ext_info_buf + ext_info_offset);
  ext_info->infoType = FWK_ADPT_EXT_SHAPE_TYPE;
  ext_info->infoLen = sizeof(int32_t);
  ext_info_offset += kExtInfoHeadSize;
  int32_t *shape_type = reinterpret_cast<int32_t *>(ext_info_buf + ext_info_offset);
  *shape_type = unknow_shape_type;
  ext_info_offset += ext_info->infoLen;
  // init ext info 2: BitMap
  ext_info = reinterpret_cast<ExtInfo *>(ext_info_buf + ext_info_offset);
  ext_info->infoType = FWK_ADPT_EXT_BITMAP;
  ext_info->infoLen = sizeof(uint64_t);
  ext_info_offset += kExtInfoHeadSize;
  uint64_t *bit_map = reinterpret_cast<uint64_t *>(ext_info_buf + ext_info_offset);
  *bit_map = 0;
  ext_info_offset += ext_info->infoLen;
  // init ext info 3: topicType
  int32_t type_value = 0;
  if (!AttrUtils::GetInt(op_desc_ptr, kTopicType, type_value)) {
    AICPUE_LOGI("Can not get attr[topic_type], op[%s], use default value[DEVICE_ONLY].",
                op_desc_ptr->GetName().c_str());
  }
  ext_info = reinterpret_cast<ExtInfo *>(ext_info_buf + ext_info_offset);
  ext_info->infoType = FWK_ADPT_EXT_TOPIC_TYPE;
  ext_info->infoLen = sizeof(FWKExtTopicType);
  ext_info_offset += kExtInfoHeadSize;
  FWKExtTopicType *topic_type = reinterpret_cast<FWKExtTopicType *>(ext_info_buf + ext_info_offset);
  *topic_type = FWKExtTopicType(type_value);
  ext_info_offset += ext_info->infoLen;
  // init ext info 4: input ShapeAndType
  ext_info = reinterpret_cast<ExtInfo *>(ext_info_buf + ext_info_offset);
  ext_info->infoType = FWK_ADPT_EXT_INPUT_SHAPE;
  ext_info->infoLen = (input_num * sizeof(ShapeAndType));
  ext_info_offset += kExtInfoHeadSize;
  ShapeAndType *inputs =
      reinterpret_cast<ShapeAndType *>(ext_info_buf + ext_info_offset);
  vector<uint32_t> inputs_type;
  vector<uint32_t> outputs_type;
  GetInOutPutsDataType(op_desc_ptr, inputs_type, outputs_type);
  for (size_t index = 0; index < inputs_type.size(); ++index) {
    inputs[index].type = inputs_type[index];
  }
  ext_info_offset += ext_info->infoLen;
  // init ext info 5: output ShapeAndType
  ext_info = reinterpret_cast<ExtInfo *>(ext_info_buf + ext_info_offset);
  ext_info->infoType = FWK_ADPT_EXT_OUTPUT_SHAPE;
  ext_info->infoLen = (output_num * sizeof(ShapeAndType));
  ext_info_offset += kExtInfoHeadSize;
  ShapeAndType *outputs =
      reinterpret_cast<ShapeAndType *>(ext_info_buf + ext_info_offset);
  for (size_t index = 0; index < outputs_type.size(); ++index) {
    outputs[index].type = outputs_type[index];
  }
  ext_info_offset += ext_info->infoLen;
  AICPUE_LOGI(
      "Make ext_info for unknown shape success: op[%s], op_type[%s], "
      "input_num[%zu], output_num[%zu], extend length[%lu].",
      op_desc_ptr->GetName().c_str(), op_desc_ptr->GetType().c_str(), input_num,
      output_num, extend_info_len);
  return SUCCESS;
}
}  // namespace aicpu
